// ======================================================================
//  cholesky.hpp  (ARCH1-only, v15: rsqrt 路径降拍 + 稳ECP)
//  变更要点：
//   - r0 / r1 / rcp_f 使用 hls::rsqrtf()，通常比 sqrt+1/x 更短、易被流水化
//   - diag 仍用 sqrtf() 保持数值一致性
//   - 继续使用 float 域缩放，避免 ap_fixed 量化/饱和
//   - 复数乘法绑定 DSP，ADD 保持 Auto，不提高关键路径
// ======================================================================
#ifndef XF_SOLVER_L1_CHOLESKY_HPP
#define XF_SOLVER_L1_CHOLESKY_HPP

#include <hls_stream.h>
#include <hls_math.h>
#include "ap_fixed.h"
#include "hls_x_complex.h"

namespace xf {
namespace solver {

// --------------------------- Helpers ---------------------------
template <typename T>
struct HlsHelper {
    typedef T real_t;
    static T zero() { return (T)0; }
    static T from_real(const real_t& r) { return r; }
    static real_t real(const T& v) { return v; }
    static real_t imag(const T&) { return (real_t)0; }
    static T make(const real_t& re, const real_t&) { return re; }
};
template <typename R>
struct HlsHelper<hls::x_complex<R> > {
    typedef R real_t;
    static hls::x_complex<R> zero() { return hls::x_complex<R>((R)0, (R)0); }
    static hls::x_complex<R> from_real(const R& r) { return hls::x_complex<R>(r, (R)0); }
    static R real(const hls::x_complex<R>& v) { return v.real(); }
    static R imag(const hls::x_complex<R>& v) { return v.imag(); }
    static hls::x_complex<R> make(const R& re, const R& im) { return hls::x_complex<R>(re, im); }
};
template <typename T> inline T conj_if_needed(const T& x) { return x; }
template <typename R> inline hls::x_complex<R> conj_if_needed(const hls::x_complex<R>& x) { return hls::x_conj(x); }

template <typename OutT>
inline typename HlsHelper<OutT>::real_t norm2(const OutT& v) {
#pragma HLS INLINE
    typedef typename HlsHelper<OutT>::real_t real_t;
    real_t re = HlsHelper<OutT>::real(v);
    real_t im = HlsHelper<OutT>::imag(v);
    real_t s0 = (real_t)(re * re);
    real_t s1 = (real_t)(im * im);
    return (real_t)(s0 + s1);
}

// 在 float 域缩放，避免 ap_fixed 溢出
template <typename T>
inline T mul_by_scalar_fp(const T& a, float s) {
#pragma HLS INLINE
    float ar = (float)a;
    float pr = ar * s;
    return (T)pr;
}
template <typename R>
inline hls::x_complex<R> mul_by_scalar_fp(const hls::x_complex<R>& a, float s) {
#pragma HLS INLINE
    float re = (float)a.real();
    float im = (float)a.imag();
    float r2 = re * s;
    float i2 = im * s;
    return hls::x_complex<R>((R)r2, (R)i2);
}

// --------------------------- Traits (ARCH=1) ---------------------------
template <bool LowerTriangularL, int RowsColsA, typename InputType, typename OutputType>
struct choleskyTraits {
    static const int ARCH = 1;
    static const int J_UNROLL = (RowsColsA < 4 ? RowsColsA : 4);
    static const int I_UNROLL = (RowsColsA < 2 ? RowsColsA : 2);
    static const int COMPLETE_REG_THRESHOLD = 12;

    typedef OutputType L_OUTPUT_T;
    typedef OutputType ACCUM_T;
    typedef typename HlsHelper<OutputType>::real_t REAL_T;
    typedef REAL_T DIAG_T;
    typedef REAL_T RECIP_DIAG_T;
};

// --------------------------- N=3 specialization (split into k0/k1/k2) ---------------------------
template <typename OutT>
struct K0_out { OutT l00, l10, l20; };
template <typename OutT>
struct K1_out { OutT l11; float r1; };

template <typename OutT>
K0_out<OutT> step_k0(OutT a00, OutT a10, OutT a20) {
#pragma HLS INLINE off
#pragma HLS bind_op op=mul impl=DSP
#pragma HLS bind_op op=add impl=Auto
    typedef typename HlsHelper<OutT>::real_t real_t;
    real_t a00r = HlsHelper<OutT>::real(a00);
    float  diag0f = hls::sqrtf((float)a00r);
    float  r0     = (a00r > (real_t)0) ? hls::rsqrtf((float)a00r) : 0.0f;

    K0_out<OutT> o;
    o.l00 = HlsHelper<OutT>::from_real((real_t)diag0f);
    o.l10 = mul_by_scalar_fp(a10, r0);
    o.l20 = mul_by_scalar_fp(a20, r0);
    return o;
}

template <typename OutT>
K1_out<OutT> step_k1(OutT a11, OutT l10) {
#pragma HLS INLINE off
#pragma HLS bind_op op=mul impl=DSP
#pragma HLS bind_op op=add impl=Auto
    typedef typename HlsHelper<OutT>::real_t real_t;
    real_t n10 = xf::solver::norm2<OutT>(l10);
    real_t a11r = HlsHelper<OutT>::real(a11);
    real_t d1   = (real_t)(a11r - n10);
    float  diag1f = hls::sqrtf((float)d1);
    float  r1     = (d1 > (real_t)0) ? hls::rsqrtf((float)d1) : 0.0f;

    K1_out<OutT> o;
    o.l11 = HlsHelper<OutT>::from_real((typename HlsHelper<OutT>::real_t)diag1f);
    o.r1  = r1; // 以 float 返回，避免量化误差/溢出
    return o;
}

template <typename OutT>
OutT step_k2(OutT a22, OutT l20, OutT l21) {
#pragma HLS INLINE off
#pragma HLS bind_op op=mul impl=DSP
#pragma HLS bind_op op=add impl=Auto
    typedef typename HlsHelper<OutT>::real_t real_t;
    real_t s20 = xf::solver::norm2<OutT>(l20);
    real_t s21 = xf::solver::norm2<OutT>(l21);
    real_t a22r= HlsHelper<OutT>::real(a22);
    real_t d2  = (real_t)(a22r - s20 - s21);
    float  diag2f = hls::sqrtf((float)d2);
    return HlsHelper<OutT>::from_real((typename HlsHelper<OutT>::real_t)diag2f);
}

template <bool LowerTriangularL, typename Traits, typename InT, typename OutT>
int choleskyCore_small_N3(hls::stream<InT>& A_in, hls::stream<OutT>& L_out) {
#pragma HLS INLINE off
#pragma HLS bind_op op=mul impl=DSP
#pragma HLS bind_op op=add impl=Auto
    // load 3x3
    OutT a00 = (OutT)A_in.read();
    OutT a01 = (OutT)A_in.read();
    OutT a02 = (OutT)A_in.read();
    OutT a10 = (OutT)A_in.read();
    OutT a11 = (OutT)A_in.read();
    OutT a12 = (OutT)A_in.read();
    OutT a20 = (OutT)A_in.read();
    OutT a21 = (OutT)A_in.read();
    OutT a22 = (OutT)A_in.read();

    // k0
    K0_out<OutT> k0 = step_k0<OutT>(a00, a10, a20);
    OutT l00 = k0.l00, l10 = k0.l10, l20 = k0.l20;

    // k1
    K1_out<OutT> k1  = step_k1<OutT>(a11, l10);
    OutT l11 = k1.l11;
    float r1 = k1.r1;

    // l21 = (a21 - l20*conj(l10)) / l11
    OutT l10c = conj_if_needed(l10);
    OutT m20  = (OutT)(l20 * l10c);
    OutT s21  = (OutT)(a21 - m20);
    OutT l21  = mul_by_scalar_fp(s21, r1);

    // k2
    OutT l22 = step_k2<OutT>(a22, l20, l21);

    // store
#pragma HLS PIPELINE II=1
    if (LowerTriangularL) {
        L_out.write(l00);                      L_out.write(HlsHelper<OutT>::zero()); L_out.write(HlsHelper<OutT>::zero());
        L_out.write(l10);                      L_out.write(l11);                     L_out.write(HlsHelper<OutT>::zero());
        L_out.write(l20);                      L_out.write(l21);                     L_out.write(l22);
    } else {
        L_out.write(conj_if_needed(l00));      L_out.write(conj_if_needed(l10));     L_out.write(conj_if_needed(l20));
        L_out.write(HlsHelper<OutT>::zero());  L_out.write(conj_if_needed(l11));     L_out.write(conj_if_needed(l21));
        L_out.write(HlsHelper<OutT>::zero());  L_out.write(HlsHelper<OutT>::zero()); L_out.write(conj_if_needed(l22));
    }
    return 0;
}

// --------------------------- Small-N generic (≤12) ---------------------------
template <bool LowerTriangularL, int N, typename Traits, typename InT, typename OutT>
int choleskyCore_small_generic(hls::stream<InT>& A_in, hls::stream<OutT>& L_out) {
#pragma HLS INLINE off
    typedef typename HlsHelper<OutT>::real_t real_t;
    const int JU = Traits::J_UNROLL;
    const int IU = Traits::I_UNROLL;

    OutT A_local[N][N];
    OutT T_local[N][N];
#pragma HLS ARRAY_PARTITION variable=A_local complete dim=1
#pragma HLS ARRAY_PARTITION variable=A_local complete dim=2
#pragma HLS ARRAY_PARTITION variable=T_local complete dim=1
#pragma HLS ARRAY_PARTITION variable=T_local complete dim=2

#pragma HLS bind_op op=mul impl=DSP
#pragma HLS bind_op op=add impl=Auto

// Load
load_rows:
    for (int i = 0; i < N; ++i) {
#pragma HLS PIPELINE II=1
    load_cols:
        for (int j = 0; j < N; ++j) {
#pragma HLS UNROLL factor=JU
            A_local[i][j] = (OutT)A_in.read();
        }
    }

// Compute
compute_k:
    for (int k = 0; k < N; ++k) {
        real_t diag_sum = (real_t)0;
    diag_block:
        for (int jb = 0; jb < k; jb += JU) {
#pragma HLS PIPELINE II=1
            real_t ps = (real_t)0;
        diag_terms:
            for (int t = 0; t < JU; ++t) {
#pragma HLS UNROLL
                int j = jb + t;
                if (j < k) ps += xf::solver::norm2<OutT>(T_local[k][j]);
            }
            diag_sum += ps;
        }

        real_t akk_re = HlsHelper<OutT>::real(A_local[k][k]);
        real_t d = (akk_re > diag_sum) ? (real_t)(akk_re - diag_sum) : (real_t)0;
        float  diag_f = hls::sqrtf((float)d);
        T_local[k][k] = HlsHelper<OutT>::from_real((real_t)diag_f);
        float  rcp_f  = (d > (real_t)0) ? hls::rsqrtf((float)d) : 0.0f;

    lower_outer_blk:
        for (int ib = k + 1; ib < N; ib += IU) {
#pragma HLS PIPELINE II=1
            OutT sum_blk[Traits::I_UNROLL];
#pragma HLS ARRAY_PARTITION variable=sum_blk complete
        init_blk:
            for (int ti = 0; ti < IU; ++ti) {
#pragma HLS UNROLL
                sum_blk[ti] = HlsHelper<OutT>::zero();
            }
        reduce_j_blk:
            for (int jb = 0; jb < k; jb += JU) {
#pragma HLS PIPELINE II=1
            reduce_j_terms:
                for (int t = 0; t < JU; ++t) {
#pragma HLS UNROLL
                    int j = jb + t;
                    if (j < k) {
                        OutT lkjc = conj_if_needed(T_local[k][j]);
                    update_rows:
                        for (int ti = 0; ti < IU; ++ti) {
#pragma HLS UNROLL
                            int i = ib + ti;
                            if (i < N) {
                                OutT lik = T_local[i][j];
                                sum_blk[ti] = sum_blk[ti] + (lik * lkjc);
                            }
                        }
                    }
                }
            }
        write_blk:
            for (int ti = 0; ti < IU; ++ti) {
#pragma HLS UNROLL
                int i = ib + ti;
                if (i < N) {
                    OutT numer = A_local[i][k] - sum_blk[ti];
                    T_local[i][k] = mul_by_scalar_fp(numer, rcp_f);
                }
            }
        }
    }

// Store
store_rows:
    for (int i = 0; i < N; ++i) {
    store_cols:
        for (int j = 0; j < N; ++j) {
#pragma HLS PIPELINE II=1
            if (LowerTriangularL) {
                L_out.write((j <= i) ? T_local[i][j] : HlsHelper<OutT>::zero());
            } else {
                if (j >= i) {
                    OutT v = T_local[j][i];
                    L_out.write(conj_if_needed(v));
                } else {
                    L_out.write(HlsHelper<OutT>::zero());
                }
            }
        }
    }
    return 0;
}

// --------------------------- Large-N (>12) ---------------------------
template <bool LowerTriangularL, int N, typename Traits, typename InT, typename OutT>
int choleskyCore_large(hls::stream<InT>& A_in, hls::stream<OutT>& L_out) {
#pragma HLS INLINE off
    typedef typename HlsHelper<OutT>::real_t real_t;
    const int JU = Traits::J_UNROLL;
    const int IU = Traits::I_UNROLL;

    OutT A_local[N][N];
    OutT T_local[N][N];
#pragma HLS ARRAY_PARTITION variable=A_local cyclic factor=Traits::J_UNROLL dim=2
#pragma HLS ARRAY_PARTITION variable=A_local cyclic factor=Traits::I_UNROLL dim=1
#pragma HLS ARRAY_PARTITION variable=T_local cyclic factor=Traits::J_UNROLL dim=2
#pragma HLS ARRAY_PARTITION variable=T_local cyclic factor=Traits::I_UNROLL dim=1
#pragma HLS bind_storage variable=T_local type=ram_2p impl=bram
#pragma HLS bind_storage variable=A_local type=ram_2p impl=bram

#pragma HLS bind_op op=mul impl=DSP
#pragma HLS bind_op op=add impl=Auto

// Load
load_rows:
    for (int i = 0; i < N; ++i) {
#pragma HLS PIPELINE II=1
    load_cols:
        for (int j = 0; j < N; ++j) {
#pragma HLS UNROLL factor=JU
            A_local[i][j] = (OutT)A_in.read();
        }
    }

// Compute
compute_k:
    for (int k = 0; k < N; ++k) {
        real_t diag_sum = (real_t)0;
    diag_block:
        for (int jb = 0; jb < k; jb += JU) {
#pragma HLS PIPELINE II=1
            real_t ps = (real_t)0;
        diag_terms:
            for (int t = 0; t < JU; ++t) {
#pragma HLS UNROLL
                int j = jb + t;
                if (j < k) ps += xf::solver::norm2<OutT>(T_local[k][j]);
            }
            diag_sum += ps;
        }

        real_t akk_re = HlsHelper<OutT>::real(A_local[k][k]);
        real_t d = (akk_re > diag_sum) ? (real_t)(akk_re - diag_sum) : (real_t)0;
        float  diag_f = hls::sqrtf((float)d);
        T_local[k][k] = HlsHelper<OutT>::from_real((real_t)diag_f);
        float  rcp_f  = (d > (real_t)0) ? hls::rsqrtf((float)d) : 0.0f;

    lower_outer_blk:
        for (int ib = k + 1; ib < N; ib += IU) {
#pragma HLS PIPELINE II=1
            OutT sum_blk[Traits::I_UNROLL];
#pragma HLS ARRAY_PARTITION variable=sum_blk complete
        init_blk:
            for (int ti = 0; ti < IU; ++ti) {
#pragma HLS UNROLL
                sum_blk[ti] = HlsHelper<OutT>::zero();
            }
        reduce_j_blk:
            for (int jb = 0; jb < k; jb += JU) {
#pragma HLS PIPELINE II=1
            reduce_j_terms:
                for (int t = 0; t < JU; ++t) {
#pragma HLS UNROLL
                    int j = jb + t;
                    if (j < k) {
                        OutT lkjc = conj_if_needed(T_local[k][j]);
                    update_rows:
                        for (int ti = 0; ti < IU; ++ti) {
#pragma HLS UNROLL
                            int i = ib + ti;
                            if (i < N) {
                                OutT lik = T_local[i][j];
                                sum_blk[ti] = sum_blk[ti] + (lik * lkjc);
                            }
                        }
                    }
                }
            }
        write_blk:
            for (int ti = 0; ti < IU; ++ti) {
#pragma HLS UNROLL
                int i = ib + ti;
                if (i < N) {
                    OutT numer = A_local[i][k] - sum_blk[ti];
                    T_local[i][k] = mul_by_scalar_fp(numer, rcp_f);
                }
            }
        }
    }

// Store
store_rows:
    for (int i = 0; i < N; ++i) {
    store_cols:
        for (int j = 0; j < N; ++j) {
#pragma HLS PIPELINE II=1
            if (LowerTriangularL) {
                L_out.write((j <= i) ? T_local[i][j] : HlsHelper<OutT>::zero());
            } else {
                if (j >= i) {
                    OutT v = T_local[j][i];
                    L_out.write(conj_if_needed(v));
                } else {
                    L_out.write(HlsHelper<OutT>::zero());
                }
            }
        }
    }
    return 0;
}

// --------------------------- Dispatcher & Public APIs ---------------------------
template <bool Small>
struct SizeDispatcher {
    template <bool LowerTriangularL, int N, typename Traits, typename InT, typename OutT>
    static int run(hls::stream<InT>& A_in, hls::stream<OutT>& L_out) {
        return choleskyCore_large<LowerTriangularL, N, Traits, InT, OutT>(A_in, L_out);
    }
};
template <>
struct SizeDispatcher<true> {
    template <bool LowerTriangularL, int N, typename Traits, typename InT, typename OutT>
    static int run(hls::stream<InT>& A_in, hls::stream<OutT>& L_out) {
        if (N == 3) {
            return choleskyCore_small_N3<LowerTriangularL, Traits, InT, OutT>(A_in, L_out);
        } else {
            return choleskyCore_small_generic<LowerTriangularL, N, Traits, InT, OutT>(A_in, L_out);
        }
    }
};

template <bool LowerTriangularL, int N, typename Traits, typename InT, typename OutT>
int choleskyCore(hls::stream<InT>& A_in, hls::stream<OutT>& L_out) {
#pragma HLS INLINE off
    return SizeDispatcher<(N <= choleskyTraits<LowerTriangularL,N,InT,OutT>::COMPLETE_REG_THRESHOLD)>::template run<LowerTriangularL, N, Traits, InT, OutT>(A_in, L_out);
}

template <bool LowerTriangularL, int RowsColsA, typename CholeskyTraits, typename InputType, typename OutputType>
int choleskyAlt(hls::stream<InputType>& A, hls::stream<OutputType>& L) {
#pragma HLS INLINE off
    return choleskyCore<LowerTriangularL, RowsColsA, CholeskyTraits, InputType, OutputType>(A, L);
}
template <bool LowerTriangularL, int RowsColsA, typename CholeskyTraits, typename InputType, typename OutputType>
int choleskyAlt2(hls::stream<InputType>& A, hls::stream<OutputType>& L) {
#pragma HLS INLINE off
    return choleskyCore<LowerTriangularL, RowsColsA, CholeskyTraits, InputType, OutputType>(A, L);
}

template <bool LowerTriangularL, int RowsColsA, typename TRAITS, typename InputType, typename OutputType>
int choleskyTop(hls::stream<InputType>& A, hls::stream<OutputType>& L) {
#pragma HLS INLINE off
#pragma HLS DATAFLOW
    return choleskyAlt<LowerTriangularL, RowsColsA, TRAITS, InputType, OutputType>(A, L);
}

template <bool LowerTriangularL, int RowsColsA, typename InputType, typename OutputType,
          typename TRAITS = choleskyTraits<LowerTriangularL, RowsColsA, InputType, OutputType> >
int cholesky(hls::stream<InputType>& A, hls::stream<OutputType>& L) {
#pragma HLS INLINE off
    return choleskyTop<LowerTriangularL, RowsColsA, TRAITS, InputType, OutputType>(A, L);
}

} // namespace solver
} // namespace xf

#endif // XF_SOLVER_L1_CHOLESKY_HPP
